package com.jiangsong.utils;

import lombok.Data;
import org.springframework.util.Assert;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;


@Data
public class MultiplyThreadLocal {

    private final Map<String, InheritableThreadLocal<Object>> sharedDataMap = new ConcurrentHashMap<>();

    public void put(String key, Object value) {
        Assert.notNull(key, "key不能为空");
        Assert.notNull(value, "value不能为空");
        InheritableThreadLocal<Object> objectInheritableThreadLocal = new InheritableThreadLocal<>();
        objectInheritableThreadLocal.set(value);
        sharedDataMap.put(key, objectInheritableThreadLocal);
    }

    public Object get(String key) {
        Assert.notNull(key, "key不能为空");
        InheritableThreadLocal<Object> objectInheritableThreadLocal = sharedDataMap.get(key);
        if (null == objectInheritableThreadLocal) {
            return null;
        }
        return objectInheritableThreadLocal.get();
    }

    public  <R> R get(String key, Class<R> clazz) {
        Assert.notNull(key, "key不能为空");
        Assert.notNull(clazz, "clazz不能为空");
        InheritableThreadLocal<Object> objectInheritableThreadLocal = sharedDataMap.get(key);
        if (null == objectInheritableThreadLocal) {
            return null;
        }
        Object o = objectInheritableThreadLocal.get();
        return (R) o;
    }

    public void removeAllThreadLocal() {
        Set<String> sharedDataKeys = sharedDataMap.keySet();
        for (String key : sharedDataKeys) {
            InheritableThreadLocal<Object> objectInheritableThreadLocal = sharedDataMap.get(key);
            if (null != objectInheritableThreadLocal) {
                System.out.println("thread remove之前:" + objectInheritableThreadLocal.get());

                objectInheritableThreadLocal.remove();
                System.out.println("thread remove之后:" + objectInheritableThreadLocal.get());
                sharedDataMap.remove(key);
            }
        }
    }

}
